import os
import random

import matplotlib.pyplot as plt
import torch

from model import UNet
from train import dataset

os.makedirs("dataset/visualization", exist_ok=True)


def detransform(output):
    output = (output + 1) * 255
    output = output.astype("uint8")
    return output


def main():
    model_path = "unet_model.pth"
    model = UNet()
    model.load_state_dict(torch.load(model_path))

    for i, (input_data, output_data) in enumerate(dataset.get_groups_images()):
        sample_index = random.randint(0, len(dataset))

        input_data, output_data = dataset[sample_index]

        input_tensor = input_data.unsqueeze(0)
        output_tensor = output_data.unsqueeze(0)

        output_tensor = model(input_tensor)

        input_data = detransform(input_data.numpy())

        output_data = detransform(output_data.numpy())

        predicted_data = detransform(output_tensor.squeeze(0).detach().numpy())

        # show the input image, output image and predicted image
        plt.figure(figsize=(15, 5))
        plt.subplot(1, 3, 1)
        plt.imshow(input_data[0], cmap="gray")
        plt.title("Input Image")
        plt.axis("off")

        plt.subplot(1, 3, 2)
        plt.imshow(output_data[0], cmap="gray")
        plt.title("Output Image")
        plt.axis("off")

        plt.subplot(1, 3, 3)
        plt.imshow(predicted_data.squeeze(0), cmap="gray")
        plt.title("Predicted Image")
        plt.axis("off")

        plt.savefig(os.path.join("assets/visualization", f"output_{i}.png"))


if __name__ == "__main__":
    main()
